package com.hub.esdemo.util;

import lombok.extern.slf4j.Slf4j;
import org.elasticsearch.action.ActionListener;
import org.elasticsearch.action.DocWriteResponse;
import org.elasticsearch.action.delete.DeleteRequest;
import org.elasticsearch.action.delete.DeleteResponse;
import org.elasticsearch.action.get.GetRequest;
import org.elasticsearch.action.get.GetResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.index.IndexResponse;
import org.elasticsearch.action.search.*;
import org.elasticsearch.action.support.WriteRequest;
import org.elasticsearch.action.update.UpdateRequest;
import org.elasticsearch.action.update.UpdateResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.client.core.CountRequest;
import org.elasticsearch.client.core.CountResponse;
import org.elasticsearch.core.TimeValue;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.query.RangeQueryBuilder;
import org.elasticsearch.index.reindex.BulkByScrollResponse;
import org.elasticsearch.index.reindex.DeleteByQueryRequest;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.SearchHits;
import org.elasticsearch.search.aggregations.AggregationBuilders;
import org.elasticsearch.search.aggregations.bucket.terms.Terms;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * @ author 飞天蝙蝠
 * @ date & time 2024/12/30 8:46
 */
@Slf4j
@Component
public class ElasticsearchUtil {

    private static RestHighLevelClient elasticsearchClient;


    public static final String ID = "id";

    public static final String LOCATION = "location";

    public static final String SCROLL_ID = "scrollId";

    @Autowired
    public void setElasticsearchClient(RestHighLevelClient elasticsearchClient) {
        ElasticsearchUtil.elasticsearchClient = elasticsearchClient;
    }

    /**
     * 查询数据
     *
     * @param index
     * @param searchSourceBuilder
     * @throws IOException
     */
    public static List<Map<String, Object>> search(String index, SearchSourceBuilder searchSourceBuilder) throws IOException {
        SearchResponse response = searchResponse(index, searchSourceBuilder);
        SearchHits hits = response.getHits();

        return Arrays.stream(hits.getHits()).map(ElasticsearchUtil::transformSearchHit).collect(Collectors.toList());
    }

    /**
     * 滚轴查询初始化
     *
     * @param index 索引
     * @param searchSourceBuilder 查询条件
     * @return List<Map<String, Object>> 返回查询结果键值对列表(包含scrollId)
     */
    public static List<Map<String, Object>> scrollInit(String index, SearchSourceBuilder searchSourceBuilder) throws IOException {
        // 设置scroll有效时间
        SearchRequest searchRequest = new SearchRequest(index);
        searchRequest.source(searchSourceBuilder);
        searchRequest.scroll(TimeValue.timeValueMinutes(3L));

        log.info("滚动查询初始化开始");
        SearchResponse response = elasticsearchClient.search(searchRequest, RequestOptions.DEFAULT);
        log.info("滚动查询初始化结束: scrollId=[{}], 数据总数=[{}], 耗时=[{}]", response.getScrollId(), response.getHits().getTotalHits().value, response.getTook());
        return scrollResult(response);
    }

    /**
     * scroll查询
     *
     * @param scrollId 滚轴查询id
     * @return List<Map<String, Object>> 返回查询结果键值对列表(包含scrollId)
     * @throws IOException
     */
    public static List<Map<String, Object>> scrollSearch(String scrollId) throws IOException {
        SearchScrollRequest scrollRequest = new SearchScrollRequest(scrollId);
        scrollRequest.scroll(TimeValue.timeValueMinutes(2L));
        SearchResponse response = elasticsearchClient.scroll(scrollRequest, RequestOptions.DEFAULT);
        log.info("滚动查询结束: 耗时=[{}], 数据总数=[{}], scrollId=[{}]", response.getTook(), response.getHits().getTotalHits().value, scrollId);

        return scrollResult(response);
    }

    /**
     * 清空scroll
     *
     * @param scrollId 滚轴查询id
     * @return boolean 成功返回true
     * @throws IOException
     */
    public static boolean scrollDelete(String scrollId) throws IOException {
        ClearScrollRequest request = new ClearScrollRequest();
        request.addScrollId(scrollId);

        ClearScrollResponse clearScrollResponse = elasticsearchClient.clearScroll(request, RequestOptions.DEFAULT);
        log.info("释放查询上下文, scrollId=[{}]", scrollId);
        return clearScrollResponse.isSucceeded();
    }

    /**
     * scroll查询结果处理
     *
     * @param response elasticsearch查询结果
     * @return List<Map<String, Object>> 返回查询结果键值对列表(包含scrollId)
     */
    private static List<Map<String, Object>> scrollResult(SearchResponse response) {
        SearchHits hits = response.getHits();
        List<Map<String, Object>> searchMap = Arrays.asList(hits.getHits()).stream().map(hit -> transformSearchHit(hit)).collect(Collectors.toList());

        // 添加scrollId
        HashMap<String, Object> tempMap = new HashMap<>(1);
        tempMap.put(SCROLL_ID, response.getScrollId());
        searchMap.add(tempMap);
        return searchMap;
    }

    /**
     * 查询返回结果
     *
     * @param index 索引名称
     * @param searchSourceBuilder 查询条件构造
     * @return SearchResponse elasticsearch查询结果
     * @throws IOException
     */
    private static SearchResponse searchResponse(String index, SearchSourceBuilder searchSourceBuilder) throws IOException {
        long beginTime = System.currentTimeMillis();
        SearchRequest searchRequest = new SearchRequest(index);
        searchRequest.source(searchSourceBuilder);

        SearchResponse response = elasticsearchClient.search(searchRequest, RequestOptions.DEFAULT);
        long endTime = System.currentTimeMillis();
        log.info("查询数据列表结束，耗时[{}]ms", endTime - beginTime);
        return response;
    }

    /**
     * 将SearchHit转为Map<String, Object>
     * kk
     * @param searchHit 单条查询结果
     * @return Map<String, Object>查询结果键值对，包含数据id
     */
    private static Map<String, Object> transformSearchHit(SearchHit searchHit) {
        Map<String, Object> dataMap = searchHit.getSourceAsMap();
        dataMap.put(ID, searchHit.getId());
        return dataMap;
    }

    /**
     * 统计总数
     *
     * @param index
     * @param queryBuilder
     * @return
     * @throws IOException
     */
    public static Long count(String index, QueryBuilder queryBuilder) throws IOException {
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        searchSourceBuilder.query(queryBuilder);


        return count(index, searchSourceBuilder);
    }

    /**
     * 统计索引总数
     *
     * @param index
     * @return
     */
    public static Long count(String index) throws IOException {
        long beginTime = System.currentTimeMillis();

        CountRequest countRequest = new CountRequest(index);
        CountResponse count = elasticsearchClient.count(countRequest, RequestOptions.DEFAULT);
        long endTime = System.currentTimeMillis();

        long totalNum = count.getCount();
        log.info("查询索引总数：count=[{}], 耗时[{}]ms", totalNum, endTime - beginTime);
        return totalNum;
    }

    /**
     * 统计总数
     *
     * @param index
     * @param searchSourceBuilder
     * @return
     * @throws IOException
     */
    public static Long count(String index, SearchSourceBuilder searchSourceBuilder) throws IOException {
        long beginTime = System.currentTimeMillis();

        CountRequest countRequest = new CountRequest(index).source(searchSourceBuilder);
        CountResponse count = elasticsearchClient.count(countRequest, RequestOptions.DEFAULT);
        long endTime = System.currentTimeMillis();

        long totalNum = count.getCount();
        log.info("查询总数：count=[{}], 耗时[{}]ms", totalNum, endTime - beginTime);
        return totalNum;
    }

    /**
     * 获取范围查询条件
     *
     * @param field
     * @param begin
     * @param end
     * @return
     */
    public static RangeQueryBuilder getRangeQueryBuilder(String field, Object begin, Object end) {
        RangeQueryBuilder rangeQueryBuilder = null;
        if (begin != null) {
            rangeQueryBuilder = QueryBuilders.rangeQuery(field).gte(begin);
        }
        if (end != null) {
            rangeQueryBuilder = (rangeQueryBuilder == null) ? QueryBuilders.rangeQuery(field).lte(end) : rangeQueryBuilder.lte(end);
        }
        return rangeQueryBuilder;
    }

    /**
     * filter query设置must条件
     *
     * @param filterQueryBuilder
     * @param queryBuilders
     */
    public static void filterMustQueryBuilder(BoolQueryBuilder filterQueryBuilder, QueryBuilder...queryBuilders) {
        Arrays.stream(queryBuilders).filter(Objects::nonNull).forEach(queryBuilder -> filterQueryBuilder.must().add(queryBuilder));
    }

    /**
     * filter query设置mustNot条件
     *
     * @param filterQueryBuilder
     * @param queryBuilders
     */
    public static void filterMustNotQueryBuilder(BoolQueryBuilder filterQueryBuilder, QueryBuilder...queryBuilders) {
        Arrays.stream(queryBuilders).filter(Objects::nonNull).forEach(queryBuilder -> filterQueryBuilder.mustNot().add(queryBuilder));
    }

    /**
     * filter query设置should条件
     *
     * @param filterQueryBuilder
     * @param queryBuilders
     */
    public static void filterShouldQueryBuilder(BoolQueryBuilder filterQueryBuilder, QueryBuilder...queryBuilders) {
        Arrays.stream(queryBuilders).filter(Objects::nonNull).forEach(queryBuilder -> filterQueryBuilder.should().add(queryBuilder));
        filterQueryBuilder.minimumShouldMatch(1);
    }

    /**
     * 插入数据
     *
     * @param index
     * @param id
     * @param dataMap
     * @return
     */
    public static boolean add(String index, String id, Map<String, Object> dataMap) throws IOException {
        log.info("es开始插入数据，id=[{}]", id);
        IndexRequest indexRequest = new IndexRequest(index).id(id).source(dataMap);
        IndexResponse response = elasticsearchClient.index(indexRequest, RequestOptions.DEFAULT);

        return operationValid(response, DocWriteResponse.Result.CREATED);
    }

    /**
     * 插入数据: 数据id已存在更新，没有则新增
     *
     * @param index
     * @param id
     * @param dataMap
     * @return
     */
    public static boolean addStrict(String index, String id, Map<String, Object> dataMap) throws IOException {
        log.info("es开始插入数据，id=[{}]", id);
        return update(index, id, dataMap);
    }

    /**
     * 批量插入数据
     *
     * @param index
     * @param data
     * @return
     */
    public static void addBatch(String index, Map<String, Map<String, Object>> data) {
        log.info("es开始批量插入数据");
        data.forEach((id, dataMap) -> {
            try {
                addStrict(index, id, dataMap);
                log.info("es插入数据成功，id=[{}]", id);
            } catch (IOException e) {
                log.error("es插入数据异常, id=[" + id + "]", e);
            }
        });
    }

    /**
     * 异步插入数据
     *
     * @param index
     * @param id
     * @param dataMap
     * @throws IOException
     */
    public static void addAsync(String index, String id, Map<String, Object> dataMap, ActionListener<IndexResponse> listener) throws IOException {
        log.info("es开始异步插入数据，id=[{}]", id);
        IndexRequest indexRequest = new IndexRequest(index).id(id).source(dataMap);
        elasticsearchClient.indexAsync(indexRequest, RequestOptions.DEFAULT, listener);
    }

    /**
     * 更新数据, 数据不存在则新增
     *
     * @param index
     * @param id
     * @param dataMap
     * @return
     * @throws IOException
     */
    public static boolean update(String index, String id, Map<String, Object> dataMap) throws IOException {
        log.info("es开始更新数据，id=[{}]", id);
        UpdateRequest updateRequest = new UpdateRequest(index, id).doc(dataMap).upsert(dataMap);
        //立即刷新
        updateRequest.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
        UpdateResponse response = elasticsearchClient.update(updateRequest, RequestOptions.DEFAULT);

        return operationValid(response, DocWriteResponse.Result.UPDATED) || operationValid(response, DocWriteResponse.Result.CREATED);
    }

    /**
     * 更新数据, 不损耗性能
     *
     * @param index
     * @param id
     * @param dataMap
     * @return
     * @throws IOException
     */
    public static boolean updateForImage(String index, String id, Map<String, Object> dataMap) throws IOException {
        log.info("es开始更新数据，id=[{}]", id);
        UpdateRequest updateRequest = new UpdateRequest(index, id).doc(dataMap).upsert(dataMap);
        UpdateResponse response = elasticsearchClient.update(updateRequest, RequestOptions.DEFAULT);

        return operationValid(response, DocWriteResponse.Result.UPDATED) || operationValid(response, DocWriteResponse.Result.CREATED);
    }

    /**
     * 异步更新数据, 数据不存在则新增
     *
     * @param index
     * @param id
     * @param dataMap
     * @return
     * @throws IOException
     */
    public static void updateAsync(String index, String id, Map<String, Object> dataMap, ActionListener<UpdateResponse> listener) throws IOException {
        log.info("es开始异步更新数据，id=[{}]", id);
        UpdateRequest updateRequest = new UpdateRequest(index, id).doc(dataMap).upsert(dataMap);
        elasticsearchClient.updateAsync(updateRequest, RequestOptions.DEFAULT, listener);
    }

    /**
     * 删除数据
     *
     * @param index
     * @param id
     * @return
     * @throws IOException
     */
    public static boolean delete(String index, String id) throws IOException {
        log.info("es开始删除数据，id=[{}]", id);
        DeleteRequest deleteRequest = new DeleteRequest(index, id);
        DeleteResponse response = elasticsearchClient.delete(deleteRequest, RequestOptions.DEFAULT);

        return operationValid(response, DocWriteResponse.Result.DELETED);
    }

    /**
     * 根据查询条件删除数据
     *
     * @param index
     * @param queryBuilder
     * @return
     * @throws IOException
     */
    public static List<Map<String, String>> deleteBatch(String index, QueryBuilder queryBuilder) throws IOException {
        log.info("es开始批量删除数据");
        DeleteByQueryRequest deleteByQueryRequest = new DeleteByQueryRequest(index).setQuery(queryBuilder);
        BulkByScrollResponse response = elasticsearchClient.deleteByQuery(deleteByQueryRequest, RequestOptions.DEFAULT);

        return response.getBulkFailures().stream().map(failure -> {
            Exception cause = failure.getCause();
            log.error("es批量删除数据失败：id=[{" + failure.getId() + "}]", cause);

            Map<String, String> valueMap = new HashMap<>(1);
            valueMap.put(failure.getId(), failure.getCause().getMessage());
            return valueMap;
        }).collect(Collectors.toList());
    }

    /**
     * 查询数据信息
     *
     * @param index 索引名称
     * @param id    数据id
     * @return
     * @throws IOException
     */
    public static Map<String, Object> select(String index, String id) throws IOException {
        GetRequest getRequest = new GetRequest(index, id);
        GetResponse response = elasticsearchClient.get(getRequest, RequestOptions.DEFAULT);

        return response.isSourceEmpty() ? null : response.getSource();
    }

    /**
     * 根据多个id查询数据
     *
     * @param index 索引名称
     * @param ids   id数组
     * @return
     * @throws IOException
     */
    public static List<Map<String, Object>> selectBatch(String index, String... ids) throws IOException {
        log.info("根据ids批量查询数据开始，ids=[" + ids + "]");
        long beginTime = System.currentTimeMillis();
        // 封装查询条件
        BoolQueryBuilder queryBuilder = QueryBuilders.boolQuery().filter(QueryBuilders.boolQuery().must(QueryBuilders.idsQuery().addIds(ids)));
        SearchRequest searchRequest = new SearchRequest(index).source(new SearchSourceBuilder().query(queryBuilder).from(0).size(ids.length));
        // 数据查询
        SearchResponse response = elasticsearchClient.search(searchRequest, RequestOptions.DEFAULT);
        SearchHit[] hits = response.getHits().getHits();
        long endTime = System.currentTimeMillis();

        log.info("根据ids批量查询数据，耗时=[{}]ms", (endTime - beginTime));
        return hits.length == 0 ? null : Stream.of(hits).map(hit -> {
            Map<String, Object> sourceAsMap = hit.getSourceAsMap();
            sourceAsMap.put(ID, hit.getId());
            return sourceAsMap;
        }).collect(Collectors.toList());
    }

    /**
     * 数据是否存在
     *
     * @param index 索引名称
     * @param id    数据id
     * @return
     * @throws IOException
     */
    public static boolean isExist(String index, String id) throws IOException {
        GetRequest getRequest = new GetRequest(index, id);
        return elasticsearchClient.exists(getRequest, RequestOptions.DEFAULT);
    }

    /**
     * 操作是否成功
     *
     * @param response  操作返回结果
     * @param result    操作类型
     * @return
     */
    private static boolean operationValid(DocWriteResponse response, DocWriteResponse.Result result) {
        return response != null && result.equals(response.getResult());
    }

    /**
     * 根据给定字段，对es数据进行聚合统计
     *
     * @param index      索引
     * @param sumField      求和字段
     * @return
     * @throws IOException
     */
    public static Map<String, Object> fieldHistogram(String index, String sumField) throws IOException {
        Map<String, Object> map = new HashMap<>();
        SearchResponse response = elasticsearchClient.search(new SearchRequest(index).source(new SearchSourceBuilder().query(QueryBuilders.boolQuery()).aggregation(AggregationBuilders.terms(sumField).field(sumField))), RequestOptions.DEFAULT);
        Terms aggregation = response.getAggregations().get(sumField);
        for (Terms.Bucket entry : aggregation.getBuckets()) {
            map.put((String) entry.getKey(), entry.getDocCount());
        }
        return map;
    }

}
